from vidgear.gears import VideoGear
import cv2
import numpy as np
import os
import logging
import time

class VideoFrameExtractor:
    def process(self, input_data):
        video_path = input_data['video_path']
        sample_rate = input_data.get('sample_rate', 5)
        
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"视频文件未找到: {video_path}")
        
        logging.info(f"开始处理视频文件: {video_path}")
        
        video = VideoGear(source=video_path).start()
        
        frames = []
        frame_count = 0
        frame_indices = []
        
        while True:
            frame = video.read()
            if frame is None:
                break
            if frame_count % sample_rate == 0:
                if frame is not None and frame.size > 0:
                    frames.append(frame)
                    frame_indices.append(frame_count)
            frame_count += 1
        
        video.stop()
        
        cap = cv2.VideoCapture(video_path)
        fps = cap.get(cv2.CAP_PROP_FPS)
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps if fps > 0 else 0
        cap.release()
        
        metadata = {
            'fps': fps,
            'width': width,
            'height': height,
            'total_frames': total_frames,
            'duration_seconds': duration,
            'sample_rate': sample_rate,
            'frame_indices': frame_indices
        }
        
        return {
            'frames': frames,
            'metadata': metadata
        }

    def process_folder(self, folder_path, output_folder="frames", sample_rate=5):
        results = {}
        
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
            print(f"创建输出文件夹: {output_folder}")
        
        video_extensions = ['.mp4', '.avi', '.mov', '.mkv', '.wmv']
        video_files = [f for f in os.listdir(folder_path) if any(f.lower().endswith(ext) for ext in video_extensions)]
        
        for video_file in video_files:
            video_path = os.path.join(folder_path, video_file)
            print(f"处理视频: {video_path}")
            
            try:
                result = self.process({
                    'video_path': video_path,
                    'sample_rate': sample_rate
                })
                
                # 使用时间戳生成唯一文件夹名，避免中文编码问题
                video_name = os.path.splitext(video_file)[0]
                unique_id = f"{int(time.time())}"  # 时间戳作为标识符
                video_frames_folder = os.path.join(output_folder, f"candidate_{unique_id}")
                if not os.path.exists(video_frames_folder):
                    os.makedirs(video_frames_folder)
                
                frames = result['frames']
                frame_indices = result['metadata']['frame_indices']
                saved_count = 0
                
                if not frames:
                    print(f"警告：从视频 {video_file} 中未提取到任何帧！")
                
                for i, (frame, frame_idx) in enumerate(zip(frames, frame_indices)):
                    # 使用英文文件名，避免中文
                    frame_path = os.path.join(video_frames_folder, f"frame_{frame_idx}.jpg")
                    success = cv2.imwrite(frame_path, frame)
                    if success:
                        saved_count += 1
                    else:
                        print(f"保存帧失败: {frame_path}")
                
                actual_files = [f for f in os.listdir(video_frames_folder) if f.endswith('.jpg')]
                if saved_count != len(actual_files):
                    print(f"警告：预期保存 {saved_count} 帧，实际找到 {len(actual_files)} 个文件")
                if not actual_files:
                    print(f"错误：帧文件夹 {video_frames_folder} 为空！")
                
                metadata = result['metadata']
                metadata['saved_frames_count'] = saved_count
                metadata['frames_folder'] = video_frames_folder
                metadata['original_name'] = video_name  # 保存原始文件名供后续使用
                results[video_file] = metadata
                
                print(f"成功提取并保存了 {saved_count} 帧到 {video_frames_folder}")
            except Exception as e:
                print(f"处理视频 {video_file} 时出错: {e}")
        
        return results

if __name__ == "__main__":
    extractor = VideoFrameExtractor()
    videos_folder = "videos"
    frames_folder = "frames"
    print(f"开始处理文件夹: {videos_folder}")
    results = extractor.process_folder(videos_folder, frames_folder)
    print("\n处理结果汇总:")
    for video_name, metadata in results.items():
        print(f"视频: {video_name}")
        print(f"  提取帧数: {metadata['saved_frames_count']}")
        print(f"  时长: {metadata['duration_seconds']:.2f} 秒")
        print(f"  尺寸: {metadata['width']}x{metadata['height']}")
        print(f"  帧保存位置: {metadata['frames_folder']}")
        print()